#!/usr/bin/env python3
# -*- coding: utf8 -*-
import os
import threading
import rospy
import rosbag
from sensor_msgs.msg import Image


def input_checked(prompt, cast):
    """带输入校验的输入函数"""
    while True:
        try:
            val = cast(input(prompt).strip())
            if cast is str and not val:
                raise ValueError
            if cast is float and val <= 0:
                raise ValueError
            return val
        except ValueError:
            print(">>> 非法输入，请重新输入！")


def main():
    rospy.init_node('bag_record_node', anonymous=True)

    name = input_checked("Bag file name (no .bag suffix): ", str)
    duration = input_checked("Recording duration (s): ", float)
    bag_name = f"{name}.bag"

    if os.path.exists(bag_name):
        print(f">>> 文件 {bag_name} 已存在，将被覆盖！")

    topic_list = ['/web_cam/image_rect_color']
    print("Recording starts in 3 s ...")
    for t in range(3, 0, -1):
        print(t, flush=True)
        rospy.sleep(1)

    # 初始化 bag
    bag = rosbag.Bag(bag_name, 'w', compression='bz2')
    msg_count = 0
    bag_closed = False
    bag_lock = threading.Lock()   # 🔒 加锁保护写入与关闭

    def cb(msg):
        nonlocal msg_count, bag_closed
        with bag_lock:
            if bag_closed:
                return
            bag.write(topic_list[0], msg, msg.header.stamp)
            msg_count += 1

    sub = rospy.Subscriber(topic_list[0], Image, cb, queue_size=200)

    rate = rospy.Rate(10)
    start = rospy.Time.now()
    end = start + rospy.Duration(duration)

    try:
        while rospy.Time.now() < end:
            elapsed = (rospy.Time.now() - start).to_sec()
            print(f"\r[{int(elapsed):02d}/{int(duration):02d}]  {msg_count} messages recorded ...",
                  end='', flush=True)
            rate.sleep()
    except KeyboardInterrupt:
        print("\n>>> 用户中断，停止录制 ...")

    # 🔒 关闭阶段，确保线程安全
    with bag_lock:
        bag_closed = True
    sub.unregister()
    rospy.sleep(0.3)  # 等待回调完全退出
    with bag_lock:
        bag.close()

    print(f"\n>>> 已保存到 ./{bag_name}  总消息数：{msg_count}")


if __name__ == '__main__':
    main()

